#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import os
import json
import shutil
from typing import List, Dict, Any
from pathlib import Path

import torch
from torch import Tensor
from tqdm import tqdm

from hparam import HyperParams
from scene.dataset_readers import SCENE_DATA_LOADERS
from scene.cameras import Camera
from scene.gaussian_model import MutilFreqGaussianModel, GaussianModel
from utils.camera_utils import load_camera
from utils.system_utils import mkdir_p


class Scene:

    def __init__(self, hp:HyperParams, load_iter=None, resolution_scales:List[float]=[1.0]):
        self.hp = hp
        self.resolution_scales = resolution_scales
        self.train_cameras: Dict[float, List[Camera]] = {}
        self.test_cameras:  Dict[float, List[Camera]] = {}
        
        self.background = (torch.ones if hp.white_background else torch.zeros)([3], dtype=torch.float, device='cuda')
        self.model_p = Path(hp.model_path)
        self.gaussians = MutilFreqGaussianModel(hp)
        self.load_iter = load_iter

        # Load scene info
        if os.path.exists(os.path.join(hp.source_path, 'sparse')):
            scene_info = SCENE_DATA_LOADERS['Colmap'](hp.source_path, hp.images, hp.eval)
        elif os.path.exists(os.path.join(hp.source_path, 'transforms_train.json')):
            print('Found transforms_train.json file, assuming Blender data set!')
            scene_info = SCENE_DATA_LOADERS['Blender'](hp.source_path, hp.white_background, hp.eval)
        else:
            raise TypeError(f'Could not recognize scene type for dataset {hp.source_path}')
        self.scene_info = scene_info

        # Init emebddings
        for gaussians in self.all_gaussians.values():
            gaussians.init_appearance_embed(len(scene_info.train_cameras))
            
        # Setup spatial_lr_scale
        self.cameras_extent: float = scene_info.nerf_normalization['radius']
        print('>> cameras_extent:', self.cameras_extent)
        self.gaussians.spatial_lr_scale = self.cameras_extent

        # Init gaussians
        if self.load_iter is not None and self.load_iter < 0:
            try:
                self.load_iter = max([int(dp.stem.split('_')[-1]) for dp in (self.model_p / 'point_cloud').iterdir()])
            except:
                print('>> not found saved point_cloud.ply')
                self.load_iter = None
        
        if self.load_iter is None:
            shutil.copyfile(scene_info.ply_path, os.path.join(self.model_p, 'input.ply'))
            cam_infos = [cam.to_json(id) for id, cam in enumerate(scene_info.train_cameras + scene_info.test_cameras)]
            with open(os.path.join(self.model_p, 'cameras.json'), 'w') as fh:
                json.dump(cam_infos, fh, indent=2, ensure_ascii=False)
            print('>> [gaussian] init via from_pcd')
            self.gaussians.from_pcd(scene_info.point_cloud)
        else:
            print(f'>> [gaussian] init via load_ply at iteration-{self.load_iter}')
            self.load_gaussian(self.load_iter)

        # Init cameras
        limit = self.hp.limit
        resolution = self.hp.resolution
        for res in self.resolution_scales:
            print('Loading Train Cameras')
            self.train_cameras[res] = [load_camera(self.hp, resolution, id, cam, res) for id, cam in enumerate(tqdm(self.scene_info.train_cameras)) if id < limit]
            print('Loading Test Cameras')
            self.test_cameras[res] = [load_camera(self.hp, resolution, id, cam, res) for id, cam in enumerate(tqdm(self.scene_info.test_cameras)) if id < limit]

    @property
    def model_path(self) -> Path: return Path(self.model_p)

    @classmethod
    def random_background(cls) -> Tensor:
        return torch.rand([3], dtype=torch.float, device='cuda')

    def get_train_cameras(self, scale:float=1.0) -> List[Camera]:
        return self.train_cameras[scale]

    def get_test_cameras(self, scale:float=1.0) -> List[Camera]:
        return self.test_cameras[scale]
    
    @property
    def cur_gaussian(self) -> GaussianModel:
        return self.gaussians.cur_gaussian

    @property
    def all_gaussians(self) -> Dict[int, GaussianModel]:
        return self.gaussians.gaussians

    def activate_gaussian(self, idx:int=0) -> GaussianModel:
        self.gaussians.activate_gaussian(idx)
        return self.cur_gaussian

    def save_gaussian(self, steps:int):
        base_dir = mkdir_p(self.model_p / 'point_cloud' / f'iteration_{steps}', parents=True)
        for idx, gaussians in self.all_gaussians.items():
            gaussians.save_ply(base_dir / f'point_cloud_freq_{idx}.ply')
            gaussians.save_pth(base_dir / f'model_freq_{idx}.pth')

    def load_gaussian(self, steps:int):
        base_dir = self.model_p / 'point_cloud' / f'iteration_{steps}'
        for idx, gaussians in self.all_gaussians.items():
            gaussians.load_ply(base_dir / f'point_cloud_freq_{idx}.ply')
            gaussians.load_pth(base_dir / f'model_freq_{idx}.pth')

    def save_checkpoint(self, steps:int):
        state_dict = {'steps': steps}
        for idx, gaussians in self.all_gaussians.items():
            try: state_dict[idx] = gaussians.state_dict()
            except AttributeError: print(f'>> freq-{idx} is not availbale')
        torch.save(state_dict, os.path.join(self.model_p, f'ckpt-{steps}.pth'))

    def load_checkpoint(self, path:str) -> int:
        state_dict: Dict[str, Any] = torch.load(path)
        steps = state_dict.get('steps', 0)
        for idx, gaussians in self.all_gaussians.items():
            try: gaussians.load_state_dict(state_dict[idx])
            except AttributeError: print(f'>> freq-{idx} is not availbale')
        return steps
